/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef LAYERNORM_OPERATION_H
#define LAYERNORM_OPERATION_H
#include <vector>
#include <string>
#include "atb/operation.h"
#include "acl/acl.h"
#include "utils.h"
#include "../custom_op_kernel/kernel_layer_norm.h"

namespace atb_plugin {

struct LayernormParam {
    float factor = 0.001f;
    float mfactor = -0.001f;
    float eps = 1e-5;
};

class LayernormOperation : public atb::Operation {
public:
    explicit LayernormOperation(const std::string &name, LayernormParam param);
    ~LayernormOperation() override;
    std::string GetName() const override;
    atb::Status InferShape(const atb::SVector<atb::TensorDesc> &inTensorDescs,
        atb::SVector<atb::TensorDesc> &outTensorDescs) const override;
    uint32_t GetInputNum() const override;
    uint32_t GetOutputNum() const override;
    atb::Status Setup(const atb::VariantPack &variantPack, uint64_t &workspaceSize, atb::Context *context) override;
    atb::Status Execute(const atb::VariantPack &variantPack, uint8_t *workspace, uint64_t workspaceSize,
        atb::Context* context) override;

private:
    std::string name_;
    LayernormParam param_;
    LayerNormTilingData tiling_;
};
} // namespace atb_plugin
#endif
